// Expression generation sub-modules
mod binary_ops;
mod calls;
mod field_access;
mod function_literal;
mod literals;
mod option;
mod panic;
mod rc_weak;
mod result;
mod slice_range;
mod structs;
mod type_utils;
mod unary_ops;

use crate::ast::*;
use crate::error::{CompilerError, Result};
use inkwell::values::BasicValueEnum;
use inkwell::AddressSpace;
use inkwell::IntPredicate;

use super::LLVMBackend;

impl<'ctx> LLVMBackend<'ctx> {
    /// Generate LLVM IR for an expression
    pub(super) fn generate_expression(
        &mut self,
        expr: &Expression,
    ) -> Result<BasicValueEnum<'ctx>> {
        match expr {
            Expression::Literal(lit) => self.generate_literal(lit),
            Expression::Variable(name) => {
                if let Some(var_value) = self.variable_map.get(name) {
                    // Check if this is a function parameter or local variable
                    if var_value.is_pointer_value() {
                        // This is a local variable (alloca), we need to load the value
                        let ptr = var_value.into_pointer_value();

                        // Check the variable's type to determine what to load
                        if let Some(var_type) = self.variable_type_map.get(name) {
                            match var_type {
                                Type::Reference(ref_type) => {
                                    // For reference types, we need to load the pointer first, then load the value
                                    // This implements automatic dereferencing
                                    let ptr_type =
                                        self.context.i8_type().ptr_type(AddressSpace::default());
                                    let ptr_value = self.builder.build_load(
                                        ptr_type,
                                        ptr,
                                        &format!("{}_ptr", name),
                                    )?;

                                    if let BasicValueEnum::PointerValue(actual_ptr) = ptr_value {
                                        // Now load the actual value from the referenced location
                                        let llvm_type = self.nrc_type_to_llvm_type(&Some(
                                            *ref_type.referenced_type.clone(),
                                        ))?;
                                        Ok(self
                                            .builder
                                            .build_load(llvm_type, actual_ptr, name)?
                                            .into())
                                    } else {
                                        Err(CompilerError::internal(
                                            "Expected pointer value for reference",
                                        ))
                                    }
                                }
                                Type::Function(func_type) => {
                                    let llvm_type = self.nrc_type_to_llvm_type(&Some(
                                        Type::Function(func_type.clone()),
                                    ))?;
                                    Ok(self.builder.build_load(llvm_type, ptr, name)?.into())
                                }
                                _ => {
                                    // Regular type, load using the correct type
                                    let llvm_type =
                                        self.nrc_type_to_llvm_type(&Some(var_type.clone()))?;
                                    Ok(self.builder.build_load(llvm_type, ptr, name)?.into())
                                }
                            }
                        } else {
                            // No type info, assume i32
                            let i32_type = self.context.i32_type();
                            Ok(self.builder.build_load(i32_type, ptr, name)?.into())
                        }
                    } else {
                        // This is a function parameter or constant, return directly
                        Ok(*var_value)
                    }
                } else if self.function_map.contains_key(name) {
                    self.build_function_value_reference(name)
                } else {
                    Err(CompilerError::internal(&format!(
                        "Undefined variable: {}",
                        name
                    )))
                }
            }
            Expression::Call(call) => self.generate_call_expression(call),
            Expression::Binary(binary) => self.generate_binary_expression(binary),
            Expression::Unary(unary) => self.generate_unary_expression(unary),
            Expression::Try(try_expr) => self.generate_try_expression(try_expr),
            Expression::ForceUnwrap(force_expr) => {
                self.generate_force_unwrap_expression(force_expr)
            }
            Expression::Match(match_expr) => self.generate_match_expression(match_expr),
            Expression::New(new_expr) => self.generate_new_expression(new_expr),
            Expression::Rc(rc_expr) => self.generate_rc_expression(rc_expr),
            Expression::FieldAccess(field_access) => {
                self.generate_field_access_expression(field_access)
            }
            Expression::MethodCall(method_call) => {
                self.generate_method_call_expression(method_call)
            }
            Expression::Struct(struct_expr) => self.generate_struct_expression(struct_expr),
            Expression::If(if_expr) => self.generate_if_expression(if_expr),
            Expression::Panic(panic_expr) => self.generate_panic_expression(panic_expr),
            Expression::Block(block_expr) => self.generate_block_expression(block_expr),
            Expression::Tuple(tuple_expr) => self.generate_tuple_expression(tuple_expr),
            Expression::Slice(slice_expr) => self.generate_slice_expression(slice_expr),
            Expression::IndexAccess(index_access) => {
                self.generate_index_access_expression(index_access)
            }
            Expression::SliceRange(slice_range) => {
                self.generate_slice_range_expression(slice_range)
            }
            Expression::Map(map_expr) => self.generate_map_expression(map_expr),
            Expression::Path(path_expr) => self.generate_path_expression(path_expr),
            Expression::FunctionLiteral(func_lit) => {
                self.generate_function_literal_expression(func_lit)
            }
            _ => Err(CompilerError::internal("Unsupported expression type")),
        }
    }

    /// Generate LLVM IR for a path expression (e.g., Status::Pending)
    fn generate_path_expression(&mut self, path_expr: &PathExpr) -> Result<BasicValueEnum<'ctx>> {
        // For now, we only support enum variant access (EnumName::VariantName)
        if path_expr.segments.len() != 2 {
            return Err(CompilerError::internal(&format!(
                "Path expression with {} segments is not yet supported. Only EnumName::VariantName is supported.",
                path_expr.segments.len()
            )));
        }

        let enum_name = &path_expr.segments[0];
        let variant_name = &path_expr.segments[1];

        // Special handling for built-in Option and Result types
        if enum_name == "Option" && variant_name == "None" {
            // Option::None is a special case - it creates an Option with tag=1 (None)
            // We create a generic None that matches common Option sizes
            // This will be cast to the correct type by the return statement or assignment

            let tag_type = self.context.i32_type();
            let tag_value = tag_type.const_int(1, false); // tag=1 for None

            // Create a reasonably-sized None struct that can accommodate most types
            // Using [8 x i8] data array (8 bytes, same as i64/f64/pointer)
            let data_array_type = self.context.i8_type().array_type(8);
            let data_value = data_array_type.const_zero(); // All zeros for data

            let none_struct_type = self
                .context
                .struct_type(&[tag_type.into(), data_array_type.into()], false);

            let none_value =
                none_struct_type.const_named_struct(&[tag_value.into(), data_value.into()]);
            return Ok(none_value.into());
        } else if enum_name == "Option" && variant_name == "Some" {
            // Option::Some requires an argument, so it should be a CallExpr
            return Err(CompilerError::internal(
                "Option::Some should be called as a function with an argument",
            ));
        } else if enum_name == "Result" && (variant_name == "Ok" || variant_name == "Err") {
            // Result::Ok and Result::Err are handled as CallExpr, not PathExpr
            return Err(CompilerError::internal(&format!(
                "Result::{} should be called as a function, not accessed as a path",
                variant_name
            )));
        }

        // Check if this is an enum variant access
        if self.enum_declarations.contains_key(enum_name) {
            // This is an enum variant, get the global constant
            let const_name = format!("{}_{}", enum_name, variant_name);

            // Get the global constant
            if let Some(global) = self.module.get_global(&const_name) {
                // Load the constant value
                let const_value = global.get_initializer().ok_or_else(|| {
                    CompilerError::internal(&format!(
                        "Enum variant '{}::{}' has no initializer",
                        enum_name, variant_name
                    ))
                })?;

                Ok(const_value)
            } else {
                Err(CompilerError::internal(&format!(
                    "Enum variant '{}::{}' not found",
                    enum_name, variant_name
                )))
            }
        } else {
            Err(CompilerError::internal(&format!(
                "Unknown enum type '{}' in path expression",
                enum_name
            )))
        }
    }

    fn generate_force_unwrap_expression(
        &mut self,
        force_expr: &ForceUnwrapExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        let value = self.generate_expression(&force_expr.expression)?;
        let operand_type = self
            .infer_expression_type(&force_expr.expression)
            .ok_or_else(|| {
                CompilerError::internal("Unable to determine type for `!` operator operand")
            })?;

        match operand_type {
            Type::Option(option_type) => {
                self.generate_option_unwrap(value, &option_type, &force_expr.location)
            }
            Type::Result(result_type) => {
                self.generate_result_unwrap(value, &result_type, &force_expr.location)
            }
            _ => Err(CompilerError::internal(
                "`!` operator can only be applied to Option or Result values",
            )),
        }
    }

    fn generate_try_expression(
        &mut self,
        try_expr: &TryExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        let value = self.generate_expression(&try_expr.expression)?;
        let operand_type = self
            .infer_expression_type(&try_expr.expression)
            .or_else(|| {
                self.current_function_return_type.as_ref().and_then(|ty| match ty {
                    Type::Option(_) | Type::Result(_) => Some(ty.clone()),
                    _ => None,
                })
            })
            .ok_or_else(|| {
                CompilerError::internal("Unable to determine type for `?` operator operand")
            })?;

        match operand_type {
            Type::Option(option_type) => {
                let function_option_type = self
                    .current_function_return_type
                    .as_ref()
                    .and_then(|ty| match ty {
                        Type::Option(func_opt) => Some(func_opt.clone()),
                        _ => None,
                    })
                    .ok_or_else(|| {
                        CompilerError::internal(
                            "Cannot use `?` on Option when the enclosing function does not return Option",
                        )
                    })?;
                self.build_option_try(value, &option_type, &function_option_type)
            }
            Type::Result(result_type) => {
                let function_result_type = self
                    .current_function_return_type
                    .as_ref()
                    .and_then(|ty| match ty {
                        Type::Result(func_res) => Some(func_res.clone()),
                        _ => None,
                    })
                    .ok_or_else(|| {
                        CompilerError::internal(
                            "Cannot use `?` on Result when the enclosing function does not return Result",
                        )
                    })?;
                self.build_result_try(value, &result_type, &function_result_type)
            }
            _ => Err(CompilerError::internal(
                "`?` operator can only be applied to Option or Result values",
            )),
        }
    }

    fn build_option_try(
        &mut self,
        option_value: BasicValueEnum<'ctx>,
        operand_option_type: &crate::ast::types::OptionType,
        function_option_type: &crate::ast::types::OptionType,
    ) -> Result<BasicValueEnum<'ctx>> {
        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot use `?` outside of a function"))?;

        let tag = self.extract_option_tag(option_value.clone())?;
        let is_some = self.builder.build_int_compare(
            IntPredicate::EQ,
            tag,
            self.context.i32_type().const_zero(),
            "option_try_is_some",
        )?;

        let some_block = self.context.append_basic_block(func, "option_try_some");
        let none_block = self.context.append_basic_block(func, "option_try_none");
        let continue_block = self.context.append_basic_block(func, "option_try_continue");

        self.builder
            .build_conditional_branch(is_some, some_block, none_block)?;

        // Success path - extract the Some value
        self.builder.position_at_end(some_block);
        let some_value = self.extract_option_some_value(option_value, operand_option_type)?;
        self.builder.build_unconditional_branch(continue_block)?;
        let ok_block = self
            .builder
            .get_insert_block()
            .ok_or_else(|| CompilerError::internal("Missing block for option try success"))?;

        // Failure path - return None from the enclosing function
        self.builder.position_at_end(none_block);
        self.execute_all_defer_statements()?;
        self.cleanup_stack_slices()?;
        let none_value = self.create_option_value(1, None, function_option_type)?;
        let _ = self.builder.build_return(Some(&none_value));

        // Continue after successful unwrapping
        self.builder.position_at_end(continue_block);
        let phi = self
            .builder
            .build_phi(some_value.get_type(), "option_try_value")?;
        phi.add_incoming(&[(&some_value, ok_block)]);

        Ok(phi.as_basic_value())
    }

    fn build_result_try(
        &mut self,
        result_value: BasicValueEnum<'ctx>,
        operand_result_type: &crate::ast::types::ResultType,
        function_result_type: &crate::ast::types::ResultType,
    ) -> Result<BasicValueEnum<'ctx>> {
        let func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("Cannot use `?` outside of a function"))?;

        let tag = self.extract_result_tag(result_value.clone())?;
        let is_ok = self.builder.build_int_compare(
            IntPredicate::EQ,
            tag,
            self.context.i32_type().const_zero(),
            "result_try_is_ok",
        )?;

        let ok_block_handle = self.context.append_basic_block(func, "result_try_ok");
        let err_block = self.context.append_basic_block(func, "result_try_err");
        let continue_block = self.context.append_basic_block(func, "result_try_continue");

        self.builder
            .build_conditional_branch(is_ok, ok_block_handle, err_block)?;

        // Ok branch - unwrap and continue
        self.builder.position_at_end(ok_block_handle);
        let ok_value =
            self.extract_result_ok_value(result_value.clone(), operand_result_type)?;
        self.builder.build_unconditional_branch(continue_block)?;
        let ok_block = self
            .builder
            .get_insert_block()
            .ok_or_else(|| CompilerError::internal("Missing block for result try success"))?;

        // Err branch - propagate error
        self.builder.position_at_end(err_block);
        self.execute_all_defer_statements()?;
        self.cleanup_stack_slices()?;
        let err_value =
            self.extract_result_err_value(result_value, operand_result_type)?;
        let propagated_err = self.create_result_value(
            err_value,
            function_result_type.err_type.as_ref(),
            function_result_type,
            "Err",
        )?;
        let _ = self.builder.build_return(Some(&propagated_err));

        // Continue with the unwrapped Ok value
        self.builder.position_at_end(continue_block);
        let phi = self
            .builder
            .build_phi(ok_value.get_type(), "result_try_value")?;
        phi.add_incoming(&[(&ok_value, ok_block)]);
        Ok(phi.as_basic_value())
    }
}
